import cv2
import numpy as np
import pywt

# 加载图像
image = cv2.imread('piano512.PNG', cv2.IMREAD_GRAYSCALE)
watermark = cv2.imread('lsc_sign512.jpg', cv2.IMREAD_GRAYSCALE)

# 执行DWT变换
coeffs_image = pywt.wavedec2(image, 'haar', level=2)
coeffs_watermark = pywt.wavedec2(watermark, 'haar', level=2)

# 将水印嵌入到小波系数中
alpha = 0.1  # 调整水印强度
watermarked_coeffs = [c_image + alpha * np.array(c_watermark) for c_image, c_watermark in zip(coeffs_image, coeffs_watermark)]

# 执行逆DWT变换
watermarked_image = pywt.waverec2(watermarked_coeffs, 'haar')

# 显示原始图像、水印和水印后的图像
cv2.imshow('Original Image', image)
cv2.imshow('Watermark', watermark)
cv2.imshow('Watermarked Image', watermarked_image.astype(np.uint8))
cv2.waitKey(0)
cv2.destroyAllWindows()

# 提取水印
extracted_coeffs = [np.array(c_image) - alpha * np.array(c_watermark) for c_image, c_watermark in zip(coeffs_image, coeffs_watermark)]
extracted_image = pywt.waverec2(extracted_coeffs, 'haar')

# 显示提取的水印
cv2.imshow('Extracted Watermark', extracted_image.astype(np.uint8))
cv2.waitKey(0)
cv2.destroyAllWindows()
